from flask import Flask, request, jsonify, render_template, url_for
from PIL import Image
import os
import numpy as np
import config
import face_recognition
from utils import faceUtils
import base64
import cv2
import io

app = Flask(__name__)


@app.route('/upload_face', methods=['POST'])
def upload_face_api():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part in the request'}), 400
    # 检查文件大小（例如，限制为 2MB）
    file = request.files['file']
    if file.content_length > 2 * 1024 * 1024:
        return jsonify({'error': 'File is too large'}), 400
    name = file.filename.split('.')[0]
    try:
        # 创建一个唯一的文件名
        image_path = os.path.join(config.UPLOAD_FOLDER, file.filename)
        encoding_path = os.path.join(config.UPLOAD_FOLDER, f"{name}.npy")
        # 将上传的文件保存到磁盘
        file.save(image_path)
        # 加载图像并转换为 RGB
        image = Image.open(image_path).convert('RGB')
        image_np = np.array(image)
        # 使用 face_recognition 提取人脸特征
        face_locations = face_recognition.face_locations(image_np)
        if len(face_locations) == 0:
            return jsonify({'error': 'No faces found in the image'}), 400
        face_encodings = face_recognition.face_encodings(image_np, face_locations)
        # 假设只处理图像中的第一张脸 将新上传的人脸特征和姓名存储到字典中
        face_encoding = face_encodings[0]
        faceUtils.face_data[name] = face_encoding
        # 将特征保存为 NumPy 文件
        np.save(encoding_path, face_encoding)
        return jsonify({'message': 'Face uploaded and features extracted successfully'}), 200
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/compare_faces', methods=['POST'])
def compare_faces_api():
    if 'face' not in request.files:
        return jsonify({'error': 'No face image part in the request'}), 400
    face = request.files['face']
    try:
        # 加载并处理上传的人脸图像
        image1 = Image.open(face).convert('RGB')
        image1_np = np.array(image1)
        face_locations1 = face_recognition.face_locations(image1_np)
        if len(face_locations1) == 0:
            return jsonify({'error': 'No faces found in the image'}), 400
        face_encodings1 = face_recognition.face_encodings(image1_np, face_locations1)
        face_encoding1 = face_encodings1[0]
        # 比较上传的人脸与已知人脸
        best_match, best_distance = faceUtils.compare_faces(face_encoding1)
        if best_match is None:
            return jsonify({'matched_name': 'Unknown', 'confidence': best_distance}), 200
        else:
            return jsonify({'matched_name': best_match, 'confidence': best_distance}), 200
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/list_faces', methods=['GET'])
def list_faces():
    images = os.listdir(config.UPLOAD_FOLDER)
    image_files = [{'filename': f, 'url': url_for('static', filename='uploads/' + f)}
                   for f in images if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]
    return jsonify(image_files)


@app.route('/delete_face', methods=['DELETE'])
def delete_face_api():
    try:
        data = request.get_json()
        filename = data.get('filename')
        
        if not filename:
            return jsonify({'error': 'Filename is required'}), 400
        
        # 构建文件路径
        image_path = os.path.join(config.UPLOAD_FOLDER, filename)
        
        # 检查文件是否存在
        if not os.path.exists(image_path):
            return jsonify({'error': 'File not found'}), 404
        
        # 获取文件名（不含扩展名）用于删除对应的.npy文件
        name = filename.split('.')[0]
        encoding_path = os.path.join(config.UPLOAD_FOLDER, f"{name}.npy")
        
        # 删除图片文件
        os.remove(image_path)
        
        # 删除对应的特征文件（如果存在）
        if os.path.exists(encoding_path):
            os.remove(encoding_path)
        
        # 从内存中的face_data字典中删除
        if name in faceUtils.face_data:
            del faceUtils.face_data[name]
        
        return jsonify({'message': f'Face {filename} deleted successfully'}), 200
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/detect_faces_video', methods=['POST'])
def detect_faces_video():
    try:
        data = request.get_json()
        image_data = data.get('image')
        enable_recognition = data.get('enable_recognition', True)
        enable_tracking = data.get('enable_tracking', True)
        
        print(f"收到视频检测请求: enable_recognition={enable_recognition}, enable_tracking={enable_tracking}")
        print(f"faceUtils.face_data 长度: {len(faceUtils.face_data)}")
        print(f"faceUtils.face_data 键: {list(faceUtils.face_data.keys())}")
        
        if not image_data:
            return jsonify({'error': 'No image data provided'}), 400
        
        # 解码base64图像数据
        image_data = image_data.split(',')[1]  # 移除data:image/jpeg;base64,前缀
        image_bytes = base64.b64decode(image_data)
        
        # 转换为OpenCV格式
        nparr = np.frombuffer(image_bytes, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if image is None:
            return jsonify({'error': 'Failed to decode image'}), 400
        
        print(f"图像尺寸: {image.shape}")
        
        # 转换为RGB格式用于face_recognition
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 检测人脸位置
        face_locations = face_recognition.face_locations(image_rgb)
        
        print(f"检测到 {len(face_locations)} 个人脸")
        
        if not face_locations:
            return jsonify({
                'faces_detected': 0,
                'faces': [],
                'recognitions': []
            })
        
        # 提取人脸特征
        face_encodings = face_recognition.face_encodings(image_rgb, face_locations)
        
        faces = []
        recognitions = []
        
        for i, (face_location, face_encoding) in enumerate(zip(face_locations, face_encodings)):
            top, right, bottom, left = face_location
            
            face_info = {
                'top': top,
                'right': right,
                'bottom': bottom,
                'left': left,
                'name': None,
                'confidence': 0.0
            }
            
            # 如果启用人脸识别，进行识别
            if enable_recognition and faceUtils.face_data:
                print(f"开始识别人脸 {i}...")
                best_match, best_distance = faceUtils.compare_faces(face_encoding)
                print(f"人脸 {i}: 最佳匹配={best_match}, 距离={best_distance}")
                
                if best_match and best_distance < 0.5:  # 提高阈值，更容易匹配
                    face_info['name'] = best_match
                    face_info['confidence'] = 1.0 - best_distance
                    
                    recognitions.append({
                        'name': best_match,
                        'confidence': 1.0 - best_distance,
                        'location': face_location
                    })
                    print(f"识别成功: {best_match} (置信度: {1.0 - best_distance:.3f})")
                else:
                    print(f"识别失败: 距离 {best_distance:.3f} 超过阈值 0.6")
            else:
                print(f"跳过识别: enable_recognition={enable_recognition}, face_data长度={len(faceUtils.face_data)}")
            
            faces.append(face_info)
        
        result = {
            'faces_detected': len(faces),
            'faces': faces,
            'recognitions': recognitions
        }
        
        print(f"返回结果: {result}")
        return jsonify(result)
        
    except Exception as e:
        print(f"视频检测错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500


@app.route('/')
def index():
    # 渲染 templates 文件夹中的 index.html 文件
    return render_template('index.html')


if __name__ == '__main__':
    faceUtils.load_known_face_encodings()
    app.run(debug=True)
